import xml.etree.ElementTree as ET
from utils.xml_utils import pretty_xml

def create_waylines_wpml_content(waypoints_coords, flight_params, total_distance, total_duration):
    """
    根据航点坐标和飞行参数生成 waylines.wpml 的内容 (XML字符串)
    waypoints_coords: [[lon, lat], [lon, lat], ...]
    """
    ET.register_namespace('', "http://www.opengis.net/kml/2.2")
    ET.register_namespace('wpml', "http://www.dji.com/wpmz/1.0.6")

    kml_element = ET.Element("{http://www.opengis.net/kml/2.2}kml")
    document_element = ET.SubElement(kml_element, "{http://www.opengis.net/kml/2.2}Document")

    # missionConfig (与template.kml中的类似，但样本中waylines.wpml的missionConfig更简洁)
    mission_config = ET.SubElement(document_element, "{http://www.dji.com/wpmz/1.0.6}missionConfig")
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}flyToWaylineMode").text = "safely"
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}finishAction").text = "goHome"
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}exitOnRCLost").text = "goContinue"
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}executeRCLostAction").text = "goBack"
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}takeOffSecurityHeight").text = str(flight_params['takeoff_security_height'])
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}globalTransitionalSpeed").text = str(flight_params['speed'])
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}globalRTHHeight").text = "100"

    drone_info = ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}droneInfo")
    ET.SubElement(drone_info, "{http://www.dji.com/wpmz/1.0.6}droneEnumValue").text = flight_params['drone_enum']
    ET.SubElement(drone_info, "{http://www.dji.com/wpmz/1.0.6}droneSubEnumValue").text = flight_params['drone_sub_enum']
    
    ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}waylineAvoidLimitAreaMode").text = "0"

    payload_info = ET.SubElement(mission_config, "{http://www.dji.com/wpmz/1.0.6}payloadInfo")
    ET.SubElement(payload_info, "{http://www.dji.com/wpmz/1.0.6}payloadEnumValue").text = flight_params['payload_enum']
    ET.SubElement(payload_info, "{http://www.dji.com/wpmz/1.0.6}payloadSubEnumValue").text = flight_params['payload_sub_enum']
    ET.SubElement(payload_info, "{http://www.dji.com/wpmz/1.0.6}payloadPositionIndex").text = "0"

    # Folder
    folder_element = ET.SubElement(document_element, "{http://www.opengis.net/kml/2.2}Folder")
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}templateId").text = "0"
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}executeHeightMode").text = "WGS84" 
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}waylineId").text = "0" 
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}distance").text = f"{total_distance:.2f}"
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}duration").text = f"{total_duration:.2f}"
    ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}autoFlightSpeed").text = str(flight_params['speed'])

    # startActionGroup (参照样本进行增强)
    start_action_group = ET.SubElement(folder_element, "{http://www.dji.com/wpmz/1.0.6}startActionGroup")
    action_gimbal_start = ET.SubElement(start_action_group, "{http://www.dji.com/wpmz/1.0.6}action")
    ET.SubElement(action_gimbal_start, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "0"
    ET.SubElement(action_gimbal_start, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "gimbalRotate"
    action_gimbal_start_param = ET.SubElement(action_gimbal_start, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFuncParam")
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalHeadingYawBase").text = "aircraft" # 参照样本
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRotateMode").text = "absoluteAngle" # 参照样本
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalPitchRotateEnable").text = "1"
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalPitchRotateAngle").text = str(flight_params.get('gimbal_pitch', -90.0))
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRollRotateEnable").text = "0"
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRollRotateAngle").text = "0"
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalYawRotateEnable").text = "0" # 对于朝下拍摄，通常不独立控制偏航或设为0
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalYawRotateAngle").text = "0"
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRotateTimeEnable").text = "0"
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRotateTime").text = "1" # 快速调整
    ET.SubElement(action_gimbal_start_param, "{http://www.dji.com/wpmz/1.0.6}payloadPositionIndex").text = "0"

    # Placemarks (Waypoints)
    num_waypoints = len(waypoints_coords)
    for i, coord in enumerate(waypoints_coords):
        placemark = ET.SubElement(folder_element, "{http://www.opengis.net/kml/2.2}Placemark")
        point = ET.SubElement(placemark, "{http://www.opengis.net/kml/2.2}Point")
        ET.SubElement(point, "{http://www.opengis.net/kml/2.2}coordinates").text = f"{coord[0]},{coord[1]}" # lon,lat

        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}index").text = str(i)
        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}executeHeight").text = str(flight_params['altitude'])
        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}waypointSpeed").text = str(flight_params['speed'])

        heading_param = ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingParam")
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingMode").text = "followWayline"
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingAngle").text = "0" 
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointPoiPoint").text = "0.000000,0.000000,0.000000" # 新增
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingAngleEnable").text = "1" # 新增 (或0，如果完全依赖followWayline)
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingPathMode").text = "followBadArc"
        ET.SubElement(heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointHeadingPoiIndex").text = "0" # 新增

        turn_param = ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}waypointTurnParam")
        turn_mode_text = ""
        turn_damping_dist_text = ""
        if i == 0 or (i == num_waypoints - 1 and waypoints_coords[i] == waypoints_coords[0] and num_waypoints > 1): # 首尾闭合点
            turn_mode_text = "toPointAndStopWithDiscontinuityCurvature"
            turn_damping_dist_text = str(flight_params.get('turn_damping_dist_stop', 0.0))
        else:
            turn_mode_text = "coordinateTurn"
            turn_damping_dist_text = str(flight_params.get('turn_damping_dist_coordinate', 10.0))
        
        ET.SubElement(turn_param, "{http://www.dji.com/wpmz/1.0.6}waypointTurnMode").text = turn_mode_text
        ET.SubElement(turn_param, "{http://www.dji.com/wpmz/1.0.6}waypointTurnDampingDist").text = turn_damping_dist_text

        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}useStraightLine").text = "1"

        # ActionGroup for startTimeLapse (应用在第一个航点上，作用于大部分航段)
        if i == 0 and flight_params.get("enable_photo_action", True) and num_waypoints >= 2:
            action_group_photo = ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}actionGroup")
            ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupId").text = "0" # 主拍照动作组
            ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupStartIndex").text = "0"
            # 作用到倒数第二个航点定义的航段结束 (0 -> 1, 1 -> 2, ..., N-2 -> N-1)
            ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupEndIndex").text = str(num_waypoints - 2 if num_waypoints > 1 else 0)
            ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupMode").text = "sequence"
            
            action_trigger_photo = ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}actionTrigger")
            ET.SubElement(action_trigger_photo, "{http://www.dji.com/wpmz/1.0.6}actionTriggerType").text = "betweenAdjacentPoints"

            # Action 1: gimbalAngleLock
            action_gimbal_lock = ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}action")
            ET.SubElement(action_gimbal_lock, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "0"
            ET.SubElement(action_gimbal_lock, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "gimbalAngleLock"
            
            # Action 2: gimbalRotate (ensure pitch)
            action_gimbal_photo_rotate = ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}action")
            ET.SubElement(action_gimbal_photo_rotate, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "1"
            ET.SubElement(action_gimbal_photo_rotate, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "gimbalRotate"
            action_gimbal_photo_rotate_param = ET.SubElement(action_gimbal_photo_rotate, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFuncParam")
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalHeadingYawBase").text = "aircraft"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRotateMode").text = "absoluteAngle"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalPitchRotateEnable").text = "1"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalPitchRotateAngle").text = str(flight_params.get('gimbal_pitch', -90.0))
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRollRotateEnable").text = "0"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRollRotateAngle").text = "0"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalYawRotateEnable").text = "0"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalYawRotateAngle").text = "0"
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}gimbalRotateTimeEnable").text = "0" # 立即生效
            ET.SubElement(action_gimbal_photo_rotate_param, "{http://www.dji.com/wpmz/1.0.6}payloadPositionIndex").text = "0"

            # Action 3: startTimeLapse
            action_start_timelapse = ET.SubElement(action_group_photo, "{http://www.dji.com/wpmz/1.0.6}action")
            ET.SubElement(action_start_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "2" 
            ET.SubElement(action_start_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "startTimeLapse"
            photo_param = ET.SubElement(action_start_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFuncParam")
            ET.SubElement(photo_param, "{http://www.dji.com/wpmz/1.0.6}payloadPositionIndex").text = "0"
            ET.SubElement(photo_param, "{http://www.dji.com/wpmz/1.0.6}useGlobalPayloadLensIndex").text = "0" # 通常为0，表示使用下面指定的
            ET.SubElement(photo_param, "{http://www.dji.com/wpmz/1.0.6}payloadLensIndex").text = flight_params.get('payload_lens_index', 'visable')
            ET.SubElement(photo_param, "{http://www.dji.com/wpmz/1.0.6}minShootInterval").text = str(flight_params.get('photo_interval_s', 5.0))

        # ActionGroup for stopTimeLapse (应用在最后一个航点上)
        if i == num_waypoints - 1 and flight_params.get("enable_photo_action", True) and num_waypoints >= 2:
            action_group_stop_photo = ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}actionGroup")
            ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupId").text = "1" # 停止拍照动作组
            ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupStartIndex").text = str(i)
            ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupEndIndex").text = str(i)
            ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionGroupMode").text = "sequence"

            action_trigger_stop_photo = ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionTrigger")
            ET.SubElement(action_trigger_stop_photo, "{http://www.dji.com/wpmz/1.0.6}actionTriggerType").text = "reachPoint"

            # Action 1: stopTimeLapse
            action_stop_timelapse = ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}action")
            ET.SubElement(action_stop_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "0"
            ET.SubElement(action_stop_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "stopTimeLapse"
            stop_photo_param = ET.SubElement(action_stop_timelapse, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFuncParam")
            ET.SubElement(stop_photo_param, "{http://www.dji.com/wpmz/1.0.6}payloadPositionIndex").text = "0"
            ET.SubElement(stop_photo_param, "{http://www.dji.com/wpmz/1.0.6}payloadLensIndex").text = flight_params.get('payload_lens_index', 'visable')

            # Action 2: gimbalAngleUnlock
            action_gimbal_unlock = ET.SubElement(action_group_stop_photo, "{http://www.dji.com/wpmz/1.0.6}action")
            ET.SubElement(action_gimbal_unlock, "{http://www.dji.com/wpmz/1.0.6}actionId").text = "1"
            ET.SubElement(action_gimbal_unlock, "{http://www.dji.com/wpmz/1.0.6}actionActuatorFunc").text = "gimbalAngleUnlock"

        # 新增：waypointGimbalHeadingParam, isRisky, waypointWorkType
        gimbal_heading_param = ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}waypointGimbalHeadingParam")
        ET.SubElement(gimbal_heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointGimbalPitchAngle").text = "0" # 通常为0，云台角度由action控制
        ET.SubElement(gimbal_heading_param, "{http://www.dji.com/wpmz/1.0.6}waypointGimbalYawAngle").text = "0"   # 通常为0

        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}isRisky").text = "0"
        ET.SubElement(placemark, "{http://www.dji.com/wpmz/1.0.6}waypointWorkType").text = "0"

    return pretty_xml(kml_element)